import os
import cv2
import sophon.sail as sail
import numpy as np
import time


def preprocess(frames):
    # Process a batch of frames
    batch_size = min(8, len(frames))
    batch_images = np.zeros((batch_size, 3, 32, 32), dtype=np.float32)
    
    for i, frame in enumerate(frames):
        # Resize the frame
        img = cv2.resize(src=frame, dsize=(32, 32), interpolation=cv2.INTER_LINEAR)
        # Convert BGR to RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # Normalize
        img = img.astype(np.float32)
        img -= 127.5
        img *= 1 / 127.5
        # Change from (H,W,C) to (C,H,W)
        img = img.transpose(2, 0, 1)
        batch_images[i] = img
    
    return batch_images


def postprocess(outputs):
    # outputs = outputs.detach().numpy()
    outputs = list(outputs.values())[0]
    pred_idx = outputs.argmax()
    pred_conf = outputs.max()
    return pred_idx, pred_conf


def infer(camera_id, model_path):
    # 1.load model
    model = sail.Engine(model_path, 0, sail.IOMode.SYSIO)
    graph_name = model.get_graph_names()[0]
    input_name = model.get_input_names(graph_name)[0]

    # Initialize camera
    cap = cv2.VideoCapture(camera_id)
    if not cap.isOpened():
        print(f"Error: Could not open camera {camera_id}")
        return

    try:
        while True:
            # Collect 8 frames
            frames = []
            for _ in range(8):
                ret, frame = cap.read()
                if not ret:
                    print("Error: Could not read frame")
                    continue
                frames.append(frame)
                
                # Optional: Display the frame
                cv2.imshow('Camera Feed', frame)
                
                # Break if 'q' is pressed
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    raise KeyboardInterrupt

            # Process batch of frames
            time1 = time.time()
            input_array = preprocess(frames)
            time2 = time.time()
            
            input_data = {input_name: input_array}
            outputs = model.process(graph_name, input_data)
            time3 = time.time()
            
            # Process each prediction in the batch
            batch_outputs = list(outputs.values())[0]
            for j in range(len(frames)):
                pred_idx = batch_outputs[j].argmax()
                pred_conf = batch_outputs[j].max()
                time4 = time.time()
                
                print("Frame {}: pred_idx={}, pred_conf={:.2f}".format(
                    j, pred_idx, pred_conf
                ))
                
            print("preprocess_time={:.2f}ms, inference_time={:.2f}ms".format(
                (time2 - time1) * 1000,
                (time3 - time2) * 1000
            ))

    except KeyboardInterrupt:
        print("\nStopping camera feed...")
    finally:
        # Release resources
        cap.release()
        cv2.destroyAllWindows()


if __name__ == '__main__':
    model_path = 'compilation.bmodel'
    camera_id = 0  # Try 0, 1, 2, etc. depending on your camera setup
    infer(camera_id, model_path)
